import cv2
import numpy as np
import os
import re
class MangaPanelExtractor:
    def __init__(self, min_panel_size=1000):
        self.min_panel_size = min_panel_size

    def preprocess_image(self, image):
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 使用高斯模糊减少噪声
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        
        # 使用自适应阈值处理
        binary = cv2.adaptiveThreshold(
            blurred, 255, cv2.ADAPTIVE_THRESH_MEAN_C, 
            cv2.THRESH_BINARY_INV, 15, 5
        )
        
        # 使用形态学操作去除噪声
        kernel = np.ones((3, 3), np.uint8)
        cleaned = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel, iterations=1)
        
        # 使用垂直和水平核进行膨胀，以连接分隔带
        vertical_kernel = np.ones((7, 1), np.uint8)
        horizontal_kernel = np.ones((1, 7), np.uint8)
        
        vertical = cv2.dilate(cleaned, vertical_kernel, iterations=3)
        horizontal = cv2.dilate(cleaned, horizontal_kernel, iterations=3)
        
        # 合并垂直和水平边缘
        edges = cv2.bitwise_or(vertical, horizontal)
        
        # 使用闭运算填充边缘内的空隙
        kernel = np.ones((5, 5), np.uint8)
        closed = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel, iterations=2)
        
        return closed

    def detect_panels(self, preprocessed_image):
        # 查找轮廓
        contours, hierarchy = cv2.findContours(
            preprocessed_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        
        # 获取图像总面积
        height, width = preprocessed_image.shape[:2]
        total_area = height * width
        min_panel_ratio = 0.005  # 进一步减小最小面积比例
        max_panel_ratio = 0.6    # 增加最大面积比例
        
        print(f"Found {len(contours)} contours")
        print(f"Total area: {total_area}, min area: {total_area * min_panel_ratio}, max area: {total_area * max_panel_ratio}")
        
        # 筛选合适大小的轮廓
        valid_panels = []
        for i, contour in enumerate(contours):
            area = cv2.contourArea(contour)
            x, y, w, h = cv2.boundingRect(contour)
            aspect_ratio = float(w) / h if h != 0 else 0
            rect_area = w * h
            extent = float(area) / rect_area
            
            # 计算轮廓的垂直位置比例
            vertical_ratio = float(y) / height
            
            # 打印调试信息
            if area > 1000:  # 只打印面积大于1000的轮廓信息
                print(f"Contour {i}: area={area}, aspect_ratio={aspect_ratio:.2f}, extent={extent:.2f}, vertical_ratio={vertical_ratio:.2f}")
            
            # 添加更多的过滤条件
            if (area > self.min_panel_size and 
                area > total_area * min_panel_ratio and 
                area < total_area * max_panel_ratio and
                0.1 < aspect_ratio < 5 and  # 放宽宽高比限制
                extent > 0.2):  # 进一步降低extent阈值
                valid_panels.append((x, y, w, h))
                print(f"Added panel: x={x}, y={y}, w={w}, h={h}")
        
        print(f"Found {len(valid_panels)} valid panels before merging")
        
        # 如果没有检测到任何分镜，尝试直接均匀分割
        if not valid_panels:
            margin = int(width * 0.1)  # 10%的边距
            panel_height = height // 5  # 均匀分成5份
            for i in range(5):
                y = i * panel_height
                h = panel_height
                valid_panels.append((margin, y, width - 2*margin, h))
                print(f"Created default panel: x={margin}, y={y}, w={width - 2*margin}, h={h}")
        
        # 合并重叠或非常接近的框
        merged_panels = self.merge_overlapping_panels(valid_panels)
        
        # 如果检测到的分镜数量少于预期，尝试分割大的分镜
        # if len(merged_panels) < 5:
        #     merged_panels = self.split_large_panels(merged_panels, height)
            
        #     # 如果分割后仍然不足5个，尝试均匀分割最大的分镜
        #     if len(merged_panels) < 5:
        #         merged_panels = self.force_split_panels(merged_panels, height, 5)
        
        print(f"Found {len(merged_panels)} panels after processing")
        
        # 按从上到下，从左到右排序
        merged_panels.sort(key=lambda x: (x[1], x[0]))
        return merged_panels

    def merge_overlapping_panels(self, panels):
        if not panels:
            return []
        
        def overlap_ratio(box1, box2):
            x1, y1, w1, h1 = box1
            x2, y2, w2, h2 = box2
            
            # 计算重叠区域
            x_left = max(x1, x2)
            y_top = max(y1, y2)
            x_right = min(x1 + w1, x2 + w2)
            y_bottom = min(y1 + h1, y2 + h2)
            
            if x_right < x_left or y_bottom < y_top:
                return 0.0
                
            intersection = (x_right - x_left) * (y_bottom - y_top)
            area1 = w1 * h1
            area2 = w2 * h2
            union = area1 + area2 - intersection
            return intersection / union  # 使用IoU而不是简单的重叠比
        
        merged = []
        used = set()
        
        for i, panel1 in enumerate(panels):
            if i in used:
                continue
                
            current_panel = list(panel1)
            used.add(i)
            
            # 检查与其他所有框的重叠
            for j, panel2 in enumerate(panels):
                if j in used or i == j:
                    continue
                    
                if overlap_ratio(current_panel, panel2) > 0.1:  # 使用更低的IoU阈值
                    x1, y1, w1, h1 = current_panel
                    x2, y2, w2, h2 = panel2
                    # 合并框
                    new_x = min(x1, x2)
                    new_y = min(y1, y2)
                    new_w = max(x1 + w1, x2 + w2) - new_x
                    new_h = max(y1 + h1, y2 + h2) - new_y
                    current_panel = [new_x, new_y, new_w, new_h]
                    used.add(j)
                    print(f"Merged panels {i} and {j}")
            
            merged.append(tuple(current_panel))
        
        return merged

    def split_large_panels(self, panels, image_height):
        """尝试将高度过大的分镜分割成多个"""
        result = []
        expected_height = image_height / 5  # 预期的单个分镜高度
        
        for x, y, w, h in panels:
            if h > expected_height * 1.8:  # 如果分镜高度超过预期的1.8倍
                # 计算应该分成几份
                num_splits = max(2, int(round(h / expected_height)))
                split_height = h / num_splits
                
                # 分割面板
                for i in range(num_splits):
                    new_y = y + int(i * split_height)
                    new_h = int(split_height)
                    result.append((x, new_y, w, new_h))
                    print(f"Split panel into: x={x}, y={new_y}, w={w}, h={new_h}")
            else:
                result.append((x, y, w, h))
        
        return result

    def force_split_panels(self, panels, image_height, target_count):
        """强制将面板分割成目标数量"""
        if not panels:
            return []
        
        result = []
        remaining = target_count - len(panels) + 1
        
        # 找到最大的面板
        max_panel = max(panels, key=lambda x: x[3])  # 按高度排序
        max_idx = panels.index(max_panel)
        
        # 分割最大的面板
        x, y, w, h = max_panel
        split_height = h / remaining
        
        # 添加其他面板
        for i, panel in enumerate(panels):
            if i != max_idx:
                result.append(panel)
        
        # 添加分割后的面板
        for i in range(remaining):
            new_y = y + int(i * split_height)
            new_h = int(split_height)
            result.append((x, new_y, w, new_h))
            print(f"Force split panel into: x={x}, y={new_y}, w={w}, h={new_h}")
        
        return result

    def extract_panels(self, image_path, output_dir):
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 读取图片
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"Could not read image: {image_path}")
        
        # 获取图片尺寸
        height, width = image.shape[:2]
        print(f"Image size: {width}x{height}")
        
        # 预处理图片
        preprocessed = self.preprocess_image(image)
        
        # 保存预处理后的图片用于调试
        cv2.imwrite(os.path.join(output_dir, "preprocessed.jpg"), preprocessed)
        
        # 检测分镜
        panels = self.detect_panels(preprocessed)
        
        # 在原图上绘制检测到的分镜位置
        debug_image = image.copy()
        for i, (x, y, w, h) in enumerate(panels):
            cv2.rectangle(debug_image, (x, y), (x+w, y+h), (0, 255, 0), 2)
            cv2.putText(debug_image, str(i+1), (x+10, y+30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        # 保存标注后的图片
        cv2.imwrite(os.path.join(output_dir, "detected_panels.jpg"), debug_image)
        
        # 提取并保存每个分镜
        for i, (x, y, w, h) in enumerate(panels):
            # 提取分镜
            panel = image[y:y+h, x:x+w]
            
            # 保存分镜
            output_path = os.path.join(
                output_dir, 
                f"panel_{i+1}.jpg"
            )
            cv2.imwrite(output_path, panel)
        
        return len(panels)

def natural_sort_key(s):
    """用于实现文件夹的自然排序"""
    return [int(text) if text.isdigit() else text.lower()
            for text in re.split('([0-9]+)', str(s))]

def main():
    # 创建 MangaPanelExtractor 实例
    extractor = MangaPanelExtractor(min_panel_size=3000)
    
    # 确保输入和输出目录存在
    input_dir = "pic"
    output_base_dir = "extracted_panels"
    
    if not os.path.exists(input_dir):
        os.makedirs(input_dir)
        print(f"Created input directory: {input_dir}")
    
    if not os.path.exists(output_base_dir):
        os.makedirs(output_base_dir)
        print(f"Created output directory: {output_base_dir}")
    
    # 获取所有 PNG 文件
    # 获取所有JPG文件
    image_files = [f for f in os.listdir(input_dir) if f.lower().endswith('.jpg') or f.lower().endswith('.png')]
    
    print(f"\n找到 {len(image_files)} 个图片文件")
    
    if not image_files:
        print(f"\n错误：在 {input_dir} 目录中没有找到图片文件！")
        return
    
    #进行自然排序
    image_files.sort(key=natural_sort_key)
    
    # 处理每个 PNG 文件
    for png_file in image_files:
        print(f"\nProcessing {png_file}...")
        input_path = os.path.join(input_dir, png_file)
        
        # 创建以 PNG 文件名（不含扩展名）命名的输出目录
        output_dir = os.path.join(output_base_dir, os.path.splitext(png_file)[0])
        
        try:
            # 如果输出目录已存在，先删除
            if os.path.exists(output_dir):
                import shutil
                shutil.rmtree(output_dir)
            
            # 处理图片
            num_panels = extractor.extract_panels(input_path, output_dir)
            print(f"Successfully extracted {num_panels} panels from {png_file} to {output_dir}")
            
        except Exception as e:
            print(f"Error processing {png_file}: {str(e)}")
            continue

if __name__ == "__main__":
    main() 